#include "upgma.h"


upgma::upgma(vector<string> *s, Biopool::GapFunction* g, Biopool::SubMatrix m) : sequences(*s), gf(g), sub(m)
{
	N = s->size();
	result = NULL;

	//initializes alignments matrix
	alignments.resize(N);
	for(int i = 0; i < N; i++)
		alignments[i].resize(N);

	//initializes distances matrix
	distances = new double*[N];
  for (int i = 0; i < N; i++)
    distances[i] = new double[N];

  computeAlignments();
  computeDistances();


  //creates initial clusters, each one containing a single sequence
  for(int i = 0; i < N; i++){
	  UpgmaCluster* c = new UpgmaCluster();
	  c->sequences.push_back(i);
	  c->height = 0;
	  clusters.push_back(c);
  }


  //calculates initial distances using the basic formula
  clustersDistances.resize(clusters.size());

  for(int i = 0; i < clusters.size(); i++){
		for(int j = 0; j < clusters.size(); j++){
		    if(i == j)
                clustersDistances[i].push_back(0);
            else
            	clustersDistances[i].push_back(calculateC2CDistance(*clusters[i], *clusters[j]));
		}
  }


}


upgma::~upgma()
{
}

//aligns every sequence with the others
void upgma::computeAlignments(){

		for(int i=0; i < N; i++){
			for(int j=0; j < i; j++){
				Biopool::SequenceData* ad = new Biopool::SequenceData(2, sequences[i], sequences[j]);
				Biopool::ScoringS2S* ss = new Biopool::ScoringS2S(&sub, ad, 0, 1.00);
				Biopool::NWAlign* nwAlign = new Biopool::NWAlign(ad, gf, ss);
				alignments[i][j] = nwAlign;
				alignments[j][i] = nwAlign;
			}
		}

}

//calculates the distances between each sequence
void upgma::computeDistances(){

	for(int i=0; i < N; i++){
		for(int j=0; j < i; j++){
			string sq1 = alignments[i][j]->getMatch()[0];
			string sq2 = alignments[i][j]->getMatch()[1];
			int diff = 0;
			int tot = 0;
			for(int k = 0; k < sq1.length(); k++){
				if(sq1[k] != sq2[k])
					diff++;
				tot++;
			}

			double f = (double)diff/(double)tot;
            double l = log(1.0-4.0*f/3.0);
			double distance = (-3.0/4.0)*l;

			distances[i][j] = distance;
			distances[j][i] = distance;

		}
	}
}

//Returns the root of the upgma tree
UpgmaCluster* upgma::getResult(){
	if (result == NULL){
		runUpgma();
		result = clusters[0];
	}
	return result;
}

//performs the tree building operations
void upgma::runUpgma(){
	while(clusters.size() != 1){
		performOnePass();
	}
}


//merges the two nearest clusters into one and adds it to the tree
void upgma::performOnePass(){

	//find nearest clusters
	double mindist = clustersDistances[0][1];
	int ind1 = 0, ind2 = 1;

	for(int i = 0; i < clusters.size(); i++){
		for(int j = i+1; j < clusters.size(); j++){
			double dis = clustersDistances[i][j];
			if(dis < mindist){
				mindist = dis;
				ind1 = i;
				ind2 = j;
			}
		}
	}

	//new cluster
	UpgmaCluster* newcluster = new UpgmaCluster();
	//merges the two nearest clusters into one
	for(int i = 0; i < clusters[ind1]->sequences.size(); i ++)
		newcluster->sequences.push_back(clusters[ind1]->sequences[i]);
	for(int i = 0; i < clusters[ind2]->sequences.size(); i ++)
		newcluster->sequences.push_back(clusters[ind2]->sequences[i]);
	//cluster height
	newcluster->height = mindist / 2.0;

	//store old clusters
	newcluster->left = clusters[ind1];
	newcluster->right = clusters[ind2];


    //calculates new distances (fast method)
	double c1 = newcluster->left->sequences.size();
	double c2 = newcluster->right->sequences.size();
	vector< double > v;
	for(int i = 0; i < clusters.size(); i++){
	    if(i == ind1)
            v.push_back(0);
        else
            v.push_back((clustersDistances[ind1][i] * c1 + clustersDistances[ind2][i] * c2) / (c1 + c2));
	}

    //inserts new distances
	clustersDistances[ind1] = v;
	for(int i = 0; i < clusters.size(); i++){
        clustersDistances[i][ind1] = v[i];
	}

	//removes unused position in the vector
	for(int i = 0; i < clusters.size(); i++){
			std::vector< double >::iterator Iter;
			Iter = clustersDistances[i].begin() + ind2 ;
			clustersDistances[i].erase(Iter);
	}
	vector< vector<double> >::iterator Iter2;
	Iter2 = clustersDistances.begin() + ind2 ;
	clustersDistances.erase(Iter2);

    //inserts new merged cluster
	clusters[ind1] = newcluster;

    //removes old cluster
	std::vector<UpgmaCluster*>::iterator Iter;
	Iter = clusters.begin() + ind2;
	clusters.erase(Iter);


}

double upgma::calculateC2CDistance(UpgmaCluster& c1, UpgmaCluster& c2){
	double a = 1.0 / (double)( c1.sequences.size() * c2.sequences.size() );
	double b = 0;
	for(int i = 0; i < c1.sequences.size(); i++){
		for(int j = 0; j < c2.sequences.size(); j++){
			b += distances[c1.sequences[i]][c2.sequences[j]];
		}
	}

	return a*b;
}
